Using assert statements for input validation that relies on external parameters (like command-line arguments or environment variables) is not robust for production code. Python's assert statements are typically compiled out when running with the -O (optimize) flag, meaning these critical checks would be skipped. It's better to explicitly raise PushRayImageError for invalid inputs to ensure consistent behavior across all environments.

        if self.ray_type in [RayType.RAY_ML, RayType.RAY_ML_EXTRA]:
            if self.python_version not in PYTHON_VERSIONS_RAY_ML:
                raise PushRayImageError(invalid_python_version)
            if self.platform not in PLATFORMS_RAY_ML:
                raise PushRayImageError(invalid_platform)
            if self.architecture not in ARCHITECTURES_RAY_ML:
                raise PushRayImageError(invalid_architecture)
        elif self.ray_type in [RayType.RAY_LLM, RayType.RAY_LLM_EXTRA]:
            if self.python_version not in PYTHON_VERSIONS_RAY_LLM:
                raise PushRayImageError(invalid_python_version)
            if self.platform not in PLATFORMS_RAY_LLM:
                raise PushRayImageError(invalid_platform)
            if self.architecture not in ARCHITECTURES_RAY_LLM:
                raise PushRayImageError(invalid_architecture)
        else:
            # ray or ray-extra
            if self.python_version not in PYTHON_VERSIONS_RAY:
                raise PushRayImageError(invalid_python_version)
            if self.platform not in PLATFORMS_RAY:
                raise PushRayImageError(invalid_platform)
            if self.architecture not in ARCHITECTURES_RAY:
                raise PushRayImageError(invalid_architecture)

https://docs.python.org/3/reference/simple_stmts.html#the-assert-statement

Huh TIL. I don't believe we're running in optimized mode today, but I'll use this anyways to start building that practice

Test does not actually test feature branch behavior

Low Severity

The test_feature_branch_non_pr test claims to test "non-PR feature branch" behavior per its name and docstring, but it uses make_ctx(python_version="3.12", commit="abc123def456") without overriding the branch parameter. The default branch in make_ctx is "master", so this test actually exercises master branch behavior instead of feature branch behavior. The test passes coincidentally because both code paths return [sha_tag, rayci_build_id], but feature branch handling is not actually being validated.

Misleading "Successfully pushed" log in dry run mode

Low Severity

The final log message states "Successfully pushed" even when dry_run=True (i.e., upload=False). In dry run mode, _copy_image only logs what would happen and returns without actually copying any images. The summary message at line 312-314 doesn't account for the dry run state, giving users incorrect feedback that images were pushed when they weren't.